import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
import pytorch_lightning as pl
from pytorch_forecasting import TimeSeriesDataSet, TemporalFusionTransformer
from pytorch_forecasting.data import GroupNormalizer
from pytorch_lightning.callbacks import EarlyStopping

# 设置随机种子保证可复现
pl.seed_everything(42)


# ------------ 1. 生成模拟数据 ------------
def generate_dummy_data(num_products=3, length=100):
    start_date = "2020-01-01"
    dates = pd.date_range(start_date, periods=length).tolist()

    data = []
    for product_id in range(1, num_products + 1):
        # 基础销量（带季节性）
        base = 50 + 20 * np.sin(2 * np.pi * np.arange(length) / 365)
        # 添加价格特征（随机波动）
        price = np.random.uniform(10, 20, size=length)
        # 添加促销活动（随机生成0/1）
        promotion = np.random.choice([0, 1], size=length, p=[0.9, 0.1])
        # 合成销量（加入噪声和促销影响）
        sales = base * (1 + 0.3 * promotion) + np.random.normal(0, 5, size=length)

        df = pd.DataFrame({
            "date": dates,
            "product_id": str(product_id),
            "sales": sales.clip(0),  # 销量不能为负
            "price": price,
            "promotion": promotion,
            "is_holiday": str(np.random.choice([0, 1], size=length, p=[0.95, 0.05]))  # 节假日标记
        })
        data.append(df)

    data = pd.concat(data, axis=0)
    data["time_idx"] = data.groupby("product_id").cumcount()  # 时间索引
    return data


data = generate_dummy_data()
data.reset_index(drop=True, inplace=True)
print("示例数据：")
print(data.head())

# ------------ 2. 创建时间序列数据集 ------------
max_encoder_length = 30  # 输入窗口长度
max_prediction_length = 7  # 预测未来7天

dataset = TimeSeriesDataSet(
    data=data,
    time_idx="time_idx",
    target="sales",
    group_ids=["product_id"],
    static_categoricals=["product_id"],  # 静态分类特征（产品ID）
    time_varying_known_categoricals=["is_holiday"],  # 已知的未来分类特征
    time_varying_known_reals=["time_idx", "price"],  # 已知的未来连续特征
    time_varying_unknown_reals=["sales"],  # 未知的时序特征（需要预测的）
    max_encoder_length=max_encoder_length,
    max_prediction_length=max_prediction_length,
    target_normalizer=GroupNormalizer(groups=["product_id"]),  # 按产品归一化
    add_relative_time_idx=True,
    add_target_scales=True
)

# 划分训练集和验证集（按时间顺序）
train_cutoff = data["time_idx"].max() - max_prediction_length * 2
train_data = data[data["time_idx"] <= train_cutoff]
val_data = data[data["time_idx"] > train_cutoff]

# 创建数据加载器
batch_size = 64
# 检查 train_data 是否为空，并明确 drop_last 参数
if not train_data.empty:
    train_dataloader = dataset.to_dataloader(
        train_data,
        batch_size=batch_size,
        shuffle=False,
        drop_last=True  # 明确指定
    )
else:
    train_dataloader = None
    print("训练数据为空，无法创建数据加载器。")

# 检查 val_data 是否为空，并明确 drop_last 参数
if not val_data.empty:
    val_dataloader = dataset.to_dataloader(
        val_data,
        batch_size=batch_size,
        shuffle=False,
        drop_last=True  # 明确指定
    )
else:
    val_dataloader = None
    print("验证数据为空，无法创建数据加载器。")

# ------------ 3. 训练模型 ------------
if train_dataloader is not None and val_dataloader is not None:
    early_stop_callback = EarlyStopping(
        monitor="val_loss",
        patience=5,
        verbose=True,
        mode="min"
    )

    trainer = pl.Trainer(
        max_epochs=20,
        accelerator="auto",
        enable_model_summary=True,
        callbacks=[early_stop_callback]
    )

    model = TemporalFusionTransformer.from_dataset(
        dataset,
        learning_rate=0.03,
        hidden_size=32,
        attention_head_size=2,
        dropout=0.1,
        output_size=7,  # 预测未来7天
        loss=torch.nn.MSELoss()
    )

    print("开始训练...")
    trainer.fit(
        model,
        train_dataloaders=train_dataloader,
        val_dataloaders=val_dataloader
    )
else:
    print("训练或验证数据为空，无法进行训练。")

# ------------ 4. 生成预测结果 ------------
# 加载最佳模型
best_model = TemporalFusionTransformer.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)

# 准备预测数据（取最后30天作为输入）
latest_data = data[data["time_idx"] > data["time_idx"].max() - max_encoder_length - max_prediction_length]
predictions = best_model.predict(
    dataset.filter(lambda x: (x["time_idx"] > data["time_idx"].max() - max_encoder_length - max_prediction_length)),
    mode="prediction",
    return_x=True
)

# 解码预测结果
raw_predictions, x = best_model.predict(
    val_dataloader,
    mode="raw",
    return_x=True,
    n_samples=100  # 概率预测采样次数
)

# 可视化预测结果
best_model.plot_prediction(x, raw_predictions, idx=0)  # 绘制第一个样本的预测
plt.show()

# 获取数值预测结果
predicted_sales = predictions.numpy()  # 形状为 (样本数, 预测步长)
actual_sales = x["decoder_target"][0].numpy()  # 实际值

# 转换为DataFrame输出
results = []
for i in range(len(predicted_sales)):
    product_id = x["product_id"][i].item()
    for t in range(max_prediction_length):
        results.append({
            "product_id": product_id,
            "time_step": t + 1,
            "predicted_sales": predicted_sales[i][t],
            "actual_sales": actual_sales[i][t]
        })

results_df = pd.DataFrame(results)
print("\n预测结果示例：")
print(results_df.head(7))  # 显示第一个产品未来7天的预测